
# terminal 中用这个命令运行这个文件的测试案例
# python -m transformer_code.model.part.attention

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from .mask import subsequentMask


# 多头注意力机制 输出 和 输入的维度 应该相同
# input  : (batch, seq_len, d_model)
# output : (batch, seq_len, d_model)

# query, key, value: [batch_size, heads, seq_len, d_k]

# DEBUG_ENABLE = True
DEBUG_ENABLE = False


class MultiHeadAttentionLayer(nn.Module):
    def __init__(self, headsNum: int, d_model: int, dropout: float = 0.1):
        super(MultiHeadAttentionLayer, self).__init__()
        # 头 的数量
        self.headsNum : int = headsNum
        # 每个头最终的输出维度
        self.d_k : int = d_model // headsNum

        # Q, K, V 三个变换矩阵
        self.qTrans = nn.Linear(d_model, d_model)
        self.kTrans = nn.Linear(d_model, d_model)
        self.vTrans = nn.Linear(d_model, d_model)

        # 最后的线性层负责整合多头注意力的输出
        self.linearOut = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    # 要特别注意 mask, 传入的 mask 的大小有下面几种情况
    # 1. encoder 中的 self-attention, mask shape 是 (batch_size, src_seq_len, src_seq_len)
    # 2. decoder 中的 self-attention, mask shape 是 (batch_size, tgt_seq_len, tgt_seq_len)
    # 3. decoder 中的 cross-attention, mask shape 是 (batch_size, tgt_seq_len, src_seq_len)
    def forward(
            self,
            query: torch.Tensor,         # (batch_size, query_seq_len, d_m)
            key: torch.Tensor,           # (batch_size, key_seq_len, d_m)
            value: torch.Tensor,         # (batch_size, value_seq_len, d_m)
            mask: torch.Tensor = None,   # (batch_size, query_seq_len, key_seq_len)
    ) -> torch.Tensor:

        batchesNum = query.size(0)

        # Function : 将数据进行拆解，拆解成多头注意力
        # 输入维度 : [batch, seq_len, d_model] 输出维度 : [batch, headers, seq_len, out_dim]
        # 先使用 view 将 d_moel 拆解成 headsNum * d_k 再交换中间两个维度
        dimTrans = lambda x: x.view(batchesNum, -1, self.headsNum, self.d_k).transpose(1, 2)

        # 计算 query, key, value
        # Q, K, V : [batch, headers, seq_len, out_dim]
        Q = dimTrans(self.qTrans(query))
        K = dimTrans(self.kTrans(key))
        V = dimTrans(self.vTrans(value))

        # 计算 scaled dot product attention score : res = softmax((QK^T) / sqrt(d_k)) * V
        # Q, K, V : [batch, headers, *seq_len, out_dim]
        # s       : [batch, headers, tgt_seq_len, src_seq_len]
        # result  : [batch, headers, tgt_seq_len, out_dim]
        s = torch.matmul(Q, K.transpose(-1, -2)) / math.sqrt(self.d_k)

        # 进行 mask 操作，在需要 屏蔽 的位置给上一个非常大的负数，使得最终 softmax 输出为 0
        if mask is not None:
            # 因为有一个 注意力头 的关系，所以要多加上一个维度
            # [batch, query_len, key_len] -> [batch, 1, query_len, key_len]
            mask = mask.unsqueeze(1)
            # 将 mask = 0 位置的数值全部替换为 -1 * 10^9
            # 自动将 mask 从 [batch, 1, query_len, key_len] 广播成 [batch, heads, query_len, key_len]
            # s = s.masked_fill(mask == 0, -1e9)
            mask.expand_as(s)

            # 2025-6-22 update: 现在强制规定 外部要处理好 mask 的 shape
            s = s.masked_fill(mask == 0, -1e9)

        # 因为 s 的 -1 维度是 query 对应 每个 key 的 score
        s = F.softmax(s, dim = -1)

        if DEBUG_ENABLE: # 输出 s 看看，理论上应该是一个下三角而且最后一行和为1
            print("attention score: ", s)

        # TODO: dropout 为什么要放在 attention score ?
        s = self.dropout(s)

        # 加权计算结果
        res = torch.matmul(s, V)

        # 将多头注意力的结合进行 concat
        # [batch, headers, seq_len, out_dim] -> [batch, seq_len, d_model]
        # 先交换第1，2维，然后进行 view 调整
        res = (
            res.transpose(1, 2)
            .contiguous()
            .view(batchesNum, -1, self.headsNum * self.d_k)
        )

        # 返回经过线性层整合后的结果
        return self.linearOut(res)


if __name__ == '__main__':
    # 验证  cross-attention， 也就是 query 来自 tgt，而 key, value 来自 src
    batch_size   = 2
    src_seq_len  = 16
    tgt_seq_len  = 10
    d_model      = 512
    heads_num    = 8

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(device)

    # encoder 输出作为 key 和 value
    enc_output = torch.randn(batch_size, src_seq_len, d_model).to(device)
    # decoder 当前步输出作为 query
    dec_output = torch.randn(batch_size, tgt_seq_len, d_model).to(device)

    # 注意：mask 应该是对 decoder 的 query 和 encoder 的 key 做的 mask
    # 大多数情况下 cross attention 的 mask 是 None，除非想屏蔽 encoder 中某些位置
    # 这里为了测试可以使用一个全 1 的 mask 或不传 mask（等价于全部可见）
    mask = torch.ones(batch_size, tgt_seq_len, src_seq_len).to(device)

    mha = MultiHeadAttentionLayer(heads_num, d_model).to(device)

    # 随机初始化参数
    for p in mha.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

    mha.eval()

    # 注意 query 是 decoder 的，key/value 是 encoder 的
    output = mha(dec_output, enc_output, enc_output, mask=None)

    print("Decoder (Query) shape: ", dec_output.shape)  # [batch_size, tgt_seq_len, d_model]
    print("Encoder (Key/Value) shape: ", enc_output.shape)  # [batch_size, src_seq_len, d_model]
    print("Output shape: ", output.shape)  # [batch_size, tgt_seq_len, d_model]

    print("Output: ", output)
